import os
import sys
sys.path.append(".")

import pyarma as pa

from unifuncs import *
from . import rungekutta as RK

def Euler1(u0, h, J, F):

    return RK.Euler1(u0, h, J(u0), lambda u: F(u) - J(u0) * u)

def StrehmelStrong2(u0, h, J, F):

    return RK.StrehmelStrong2(u0, h, J(u0), lambda u: F(u) - J(u0) * u)

def StrehmelWeak2(u0, h, J, F):

    return RK.StrehmelWeak2(u0, h, J(u0), lambda u: F(u) - J(u0) * u)

def Cox4(u0, h, J, F):

    return RK.Cox4(u0, h, J(u0), lambda u:F(u) - J(u0) * u)

def Krogstad4(u0, h, J, F):

    return RK.Krogstad4(u0, h, J(u0), lambda u:F(u) - J(u0) * u)

def Strehmel4(u0, h, J, F):

    return RK.Strehmel4(u0, h, J(u0), lambda u: F(u) - J(u0) * u)

def Hochbruck5(u0, h, J, F):

    return RK.Hochbruck5(u0, h, J(u0), lambda u: F(u) - J(u0) * u)

def Exp32(u0, h, J, F):
    
    Jn = J(u0)
    g = lambda u: F(u) - Jn * u
    c = [0, 1]
    Un = []
    Gn = []
    
    Un.append(u0)
    Gn.append(g(Un[0]))
    
    Un.append(u0 + c[1] * h * pv.phiv(c[1] * h * Jn, [Gn[0] + Jn * u0]))
    Gn.append(g(Un[1]))
    
    u1 = u0 + h * pv.phiv(h * Jn, [Gn[0] + Jn * u0, pa.zeros(pa.size(u0)), -2 * Gn[0] + 2 * Gn[1]])
    e1 = u0 + h * pv.phiv(h * Jn, [Gn[0] + Jn * u0])
    
    return u1, pa.abs(u1 - e1)

def Exp43(u0, h, J, F):
    
    Jn = J(u0)
    g = lambda u: F(u) - Jn * u
    c = [0, 1/2, 1]
    Un = []
    Gn = []
    
    Un.append(u0)
    Gn.append(g(Un[0]))
    
    Un.append(u0 + c[1] * h * pv.phiv(c[1] * h * Jn, [Gn[0] + Jn * u0]))
    Gn.append(g(Un[1]))
    
    Un.append(u0 + c[2] * h * pv.phiv(c[2] * h * Jn, [Gn[1] + Jn * u0]))
    Gn.append(g(Un[2]))
    
    u1 = u0 + h * pv.phiv(h * Jn, [Gn[0] + Jn * u0, pa.zeros(pa.size(u0)), -14 * Gn[0] + 16 * Gn[1] - 2 * Gn[2], 36 * Gn[0] - 48 * Gn[1] + 12 * Gn[2]])
    e1 = u0 + h * pv.phiv(h * Jn, [Gn[0] + Jn * u0, pa.zeros(pa.size(u0)), -14 * Gn[0] + 16 * Gn[1] - 2 * Gn[2]])
    
    return u1, pa.abs(u1 - e1)